import jsonlines
from tqdm import tqdm
all_datas = []

import sys
slices = int(sys.argv[1])
input_file = sys.argv[2]

with jsonlines.open(input_file) as f:
    for line in tqdm(f):
        all_datas.append(line)

idx = 0

split = [[] for i in range(slices)]
for li, line in tqdm(enumerate(all_datas)):
    split[li % slices].append(line)

for idx in range(slices):
    with jsonlines.open(f"{input_file}_{idx}","w") as wf:
        for line in tqdm(split[idx]):
            wf.write(line)
